package com.cike.sparkstudy.sql.scala

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._

class StringCounter extends UserDefinedAggregateFunction {
  //inputSchema指的是输入数据类型
  def inputSchema: StructType = {StructType(Array(StructField("str",StringType,true)))}

  //中间聚合时所处理数据的类型
  def bufferSchema: StructType = {StructType(Array(StructField("count",IntegerType,true)))}

  //函数返回值类型
  def dataType: DataType = {IntegerType}

  //
  def deterministic: Boolean = {true}

  //为每个分组的数据执行初始化操作
  def initialize(buffer: MutableAggregationBuffer): Unit = {buffer(0) = 0}

  //每个分组有新数据进来时，进行分组对应的聚合值计算
  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {buffer(0) = buffer.getAs[Int](0) +1}

  // 由于Spark是分布式的，所以一个分组的数据，可能会在不同的节点上进行局部聚合，就是update
  // 但是，最后一个分组，在各个节点上的聚合值，要进行merge，也就是合并
  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = { buffer1(0) = buffer1.getAs[Int](0)+buffer2.getAs[Int](0)}

  // 最后，指的是，一个分组的聚合值，如何通过中间的缓存聚合值，最后返回一个最终的聚合值
  def evaluate(buffer: Row): Any = { buffer.getAs[Int](0)}
}
